package com.example.websocketdemo.config;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpHeaders;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

/**
 * @author Cary
 * @date 2025-03-26
 */
@Configuration
@EnableWebSocketMessageBroker
@Slf4j
public class StompConfiguration implements WebSocketMessageBrokerConfigurer {
  /**
   * 这里定义的是客户端接收服务端消息的相关信息
   *
   * @param registry
   */
  @Override
  public void configureMessageBroker(MessageBrokerRegistry registry) {
    // 消客户端订阅消息的请求前缀，notice用于广播推送，queue用于点对点推送
    registry.enableSimpleBroker("/notice", "/queue");

    // 表示配置一个或多个前缀，通过这些前缀过滤出需要被注解方法处理的消息。
    registry.setApplicationDestinationPrefixes("/app");

    // 点对点消息前缀
    registry.setUserDestinationPrefix("/user");
  }

  /**
   * 添加一个服务端点，来接收客户端的连接
   *
   * @param registry
   */
  @Override
  public void registerStompEndpoints(StompEndpointRegistry registry) {
    registry.addEndpoint("/ws")
        .setAllowedOriginPatterns("*") // 跨域
        .addInterceptors(new AuthHandshakeInterceptor())
        .withSockJS();
  }

  /**
   * 配置客户端入站通道拦截器
   * <p>
   * 添加 ChannelInterceptor 拦截器，用于在消息发送前，从请求头中获取 token 并解析出用户信息(username)，用于点对点发送消息给指定用户
   *
   * @param registration 通道注册器
   */
  @Override
  public void configureClientInboundChannel(ChannelRegistration registration) {
    registration.interceptors(new ChannelInterceptor() {
      @Override
      public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        // 如果是连接请求（CONNECT 命令），从请求头中取出 token 并设置到认证信息中
        if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {
          // TODO 从连接头中提取授权令牌验证并提取用户信息
          String bearerToken = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION);
          // 用户信息
          String username = accessor.getFirstNativeHeader("username");

          if (StrUtil.isNotBlank(username)) {
            accessor.setUser(() -> username);
            return message;
          }
        }
        // 不是连接请求，直接放行
        return ChannelInterceptor.super.preSend(message, channel);
      }
    });
  }
}
